import pyrealsense2 as rs
import numpy as np
import cv2
from ultralytics import YOLO

# 加载YOLOv8n模型（需提前下载 yolov8n.pt 到本地目录）
model = YOLO('yolov8n.pt')

# 配置 RealSense 流
pipeline = rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
pipeline.start(config)

try:
    while True:
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        depth_frame = frames.get_depth_frame()
        if not color_frame or not depth_frame:
            continue

        color_image = np.asanyarray(color_frame.get_data())
        depth_image = np.asanyarray(depth_frame.get_data())

        # YOLOv8推理
        results = model(color_image)
        for r in results:
            boxes = r.boxes
            for box in boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                conf = float(box.conf[0])
                cls = int(box.cls[0])
                # 获取检测框中心点的深度
                cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
                depth = depth_image[cy, cx]
                depth_m = depth / 1000.0
                label = f"{model.names[cls]} {conf:.2f} {depth_m:.2f}m"
                cv2.rectangle(color_image, (x1, y1), (x2, y2), (0,255,0), 2)
                cv2.putText(color_image, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)

        cv2.imshow('YOLOv8 + RealSense', color_image)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
finally:
    pipeline.stop()
    cv2.destroyAllWindows()